import os
import time
import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import torch
import torchvision as tv
import torchvision.transforms as transforms
from args import *
from tool import *
import evaluate
# from models import *


save_path = os.path.join(args.log_path, "image.{}".format(args.model_type))
if not os.path.exists(save_path):
    os.makedirs(save_path)


test_loader = torch.utils.data.DataLoader(
    tv.datasets.MNIST(args.data_path, train=False,
                   transform=transforms.Compose([
                    #    transforms.Resize([224, 224]),
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                    #    transforms.Lambda(lambda x: x.repeat(3, 1, 1))
                   ])),
    batch_size=100, shuffle=True)


P_list = [
    "centre_0bias_0.0001lambda_20scale_1dy-scale",
    "centre_0bias_0.01lambda_20scale_1dy-scale",
    "centre_0bias_0.1lambda_20scale_1dy-scale",
    "centre_0bias_1.0lambda_20scale_1dy-scale"
]
P_list = [os.path.join(args.log_path, _p) for _p in P_list]
# N_PATH = len(P_list)

M_list = [
    "27e_0.9882acc.pth",
    "21e_0.9904acc.pth",
    "20e_0.9898acc.pth",
    "23e_0.9904acc.pth"
]
M_list = [os.path.join(_p, _m) for _p, _m in zip(P_list, M_list)]
M_list = [torch.load(_m).cuda() for _m in M_list]

C_list = ["C.{}e.npy".format(i) for i in (27, 21, 20, 23)]
C_list = [os.path.join(_p, _c) for _p, _c in zip(P_list, C_list)]
C_list = [np.load(_c) for _c in C_list]
N_CLASS = C_list[0].shape[0]


def get_fea():
    Fs_list, L = [], []
    for i in range(len(M_list)):
        Fs_list.append([])
    _cnt = 0
    with torch.no_grad():
        for X, Y in test_loader:
            # X, Y = X.cuda(), Y.cuda()
            X = X.cuda()
            for m, f_list in zip(M_list, Fs_list):
                _f, _ = m(X)
                f_list.append(_f.cpu().numpy())
            L.append(Y.numpy())
            _cnt += _f.size(0)
            if _cnt >= 2000:
                break

    for i in range(len(Fs_list)):
        Fs_list[i] = np.vstack(Fs_list[i])
    L = np.concatenate(L)
    return Fs_list, L


Fs_list, L = get_fea()
W_list = []  # [10, 2] x 4
for model in M_list:
    W_list.append(model.state_dict()["clf_layer.weight"].cpu().numpy())
title_list = ("lambda = 0.0001", "lambda = 0.01", "lambda = 0.1", "lambda = 1")


title = "feature and centre"
fig, ax = plt.subplots(2, 2)
plt.title(title)
ax_list = []
for i in range(2):
    for j in range(2):
        ax_list.append(ax[i][j])
for _ax, _t, _F, _C in zip(ax_list, title_list, Fs_list, C_list):
    _ax.set_title(_t)
    _ax.scatter(_F[:, 0], _F[:, 1],
        s=10, c=L, marker='.', cmap="rainbow")
    for c in _C:
        _ax.plot([0, c[0]], [0, c[1]])
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)


title = "feature and weight vector"
fig, ax = plt.subplots(2, 2)
plt.title(title)
ax_list = []
for i in range(2):
    for j in range(2):
        ax_list.append(ax[i][j])
for _ax, _t, _F, _W in zip(ax_list, title_list, Fs_list, W_list):
    _ax.set_title(_t)
    _ax.scatter(_F[:, 0], _F[:, 1],
        s=10, c=L, marker='.', cmap="rainbow")
    for w in _W:
        _ax.plot([0, w[0]], [0, w[1]])
fig.savefig(os.path.join(save_path, "{}.png".format(title.replace(" ", "_"))))
plt.close(fig)
